package org.example;

import org.example.util.BTreeUtil;

/**
 * 红黑树实现
 * 2022-1-28 ~ 2022-2-2 v1.0
 * @param <K>
 * @param <V>
 */
public class RedBlackTree<K extends Comparable<K>, V> implements ITree<K, V>{
    public static final boolean RED   = true;
    public static final boolean BLACK = false;

    public static class RBNode <K, V> extends BSTree.BSNode<K, V> {
        public boolean getColor() {
            return color;
        }

        private boolean color = RED;

        public RBNode(RBNode<K, V> parent, K key, V value) {
            super(parent, key, value);
        }

        @Override
        public RBNode<K, V> getParent() {
            return (RBNode<K, V>) super.getParent();
        }

        public void setParent(RBNode<K, V> parent) {
            super.setParent(parent);
        }

        @Override
        public RBNode<K, V> getLeft() {
            return (RBNode<K, V>) super.getLeft();
        }

        public void setLeft(RBNode<K, V> left) {
            super.setLeft(left);
        }

        public RBNode<K, V> getRight() {
            return (RBNode<K, V>) super.getRight();
        }

        public void setRight(RBNode<K, V> right) {
            super.setRight(right);
        }

    }

    private final BSTree<K, V> bsTree = new BSTree<>();

    public RBNode<K, V> getRoot() {
        return (RBNode<K, V>) bsTree.getRoot();
    }

    public void setRoot(RBNode<K, V> root) {
        bsTree.setRoot(root);
    }

    @Override
    public void put(K key, V value) {
        RBNode<K, V> node = new RBNode<>(null, key, value);
        if (bsTree.getRoot() == null) {
            setColor(node, BLACK);
            bsTree.setRoot(node);
            return;
        }
        bsTree.put(node);
        // 平衡
        fixAfterPut(node);
    }

    @Override
    public V get(K key) {
        return bsTree.get(key);
    }

    @Override
    public V remove(K key) {
        RBNode<K, V> target = this.getNode(key);
        if (target == null) {
            return null;
        }
        V ret = target.getValue();
        RBNode<K, V> predecessor = this.getPredecessor(target);
        if (predecessor != null) {
            target.setKey(predecessor.getKey());
            target.setValue(predecessor.getValue());
            target = predecessor;
        }
        fixAndRemove234LeaveNode(target);
        return ret;
    }

    private void doRemove(RBNode<K, V> target) {
        if (target.getParent().getLeft() == target) {
            target.getParent().setLeft(target.getLeft());
        } else {
            target.getParent().setRight(target.getLeft());
        }
        target.setParent(null);
    }

    private void fixAndRemove234LeaveNode(RBNode<K, V> _234LeaveNode) {
        if (_234LeaveNode == null) { throw new NullPointerException("_234LeaveNode could not be null"); }
        if (colorOf(_234LeaveNode) != BLACK) { // 情况0, 如果自己是红节点直接删除
            doRemove(_234LeaveNode);
            return;
        }
        // 情况1, 自己有红儿子
        RBNode<K, V> left = _234LeaveNode.getLeft();
        if (left != null) {     // 情况1, 将要删除的黑节点有红儿子,让红儿子代替自己被删
            _234LeaveNode.setKey(left.getKey());
            _234LeaveNode.setValue(left.getValue());
            doRemove(left);
            return;
        }
        RBNode<K, V> right = _234LeaveNode.getRight();
        if (right != null) {    // 情况1, 将要删除的黑节点有红儿子,让红儿子代替自己被删
            _234LeaveNode.setKey(right.getKey());
            _234LeaveNode.setValue(right.getValue());
            doRemove(right);
            return;
        }

        if (_234LeaveNode.getParent() == null) {
            setRoot(null);
            return;
        }

        RBNode<K, V> p;
        RBNode<K, V> brother;
        RBNode<K, V> self = _234LeaveNode;
        do {
            p = self.getParent();
            brother = (RBNode<K, V>) BTreeUtil.getBrother(self);

            // 情况2, 自己无红儿子, 但兄弟有红儿子
            if (colorOf(brother) != RED && (brother.getRight() != null || brother.getLeft() != null)) { // 234树的叶子节点不可能无兄弟
                // 兄弟有红儿子,借兄弟儿子,兄弟和父亲交换颜色并顶替父亲位置. 情况2
                if (self == p.getLeft()) {
                    // 兄弟只有一个左儿子且与我平行,去他妈的! 右旋兄弟让和其儿子互换位置并互换颜色
                    if (brother.getRight() == null) {
                        switchColor(brother, brother.getLeft());
                        rightRotate(brother);
                        brother = brother.getParent(); // 新的兄弟出现了, 指向其原来的左儿子.
                    }
                    switchColor(brother, p);
                    setColor(brother.getRight(), BLACK);
                    leftRotate(p);
                } else { // 兄弟有红儿子,兄弟和父亲交换颜色并顶替父亲位置. 情况2的对称版本
                    // 兄弟只有一个右儿子且与我平行,去他妈的! 左旋兄弟让和其儿子互换位置并互换颜色
                    if (brother.getLeft() == null) {
                        switchColor(brother, brother.getRight());
                        leftRotate(brother);
                        brother = brother.getParent(); // 新的兄弟出现了, 指向其原来的右儿子.
                    }
                    switchColor(brother, p);
                    setColor(brother.getLeft(), BLACK);
                    rightRotate(p);
                }
                doRemove(self);
                return;
            }
            // 情况3, 父亲是红节点, 且兄弟无红儿子
            if (colorOf(p) == RED) {
                doRemove(self);
                switchColor(p, brother);
                return;
            }
            // 情况4, 父亲是黑节点, 且兄弟是红
            if (colorOf(p) == BLACK && colorOf(brother) == RED) {
                // 旋转变色让兄弟顶替父亲位置并得到新兄弟
                setColor(brother, BLACK);
                if (self == p.getLeft()) {
                    brother = brother.getLeft();
                    leftRotate(p);
                } else {
                    brother = brother.getRight();
                    rightRotate(p);
                }
                // 如果新兄弟节点有儿子
                if (brother.getLeft() != null || brother.getRight() != null) {
                    p = self.getParent();
                    if (self == p.getLeft()) {
                        if (brother.getRight() == null) {
                            rightRotate(brother);
                            brother = p.getRight();
                        }
                        switchColor(brother, brother.getRight());
                    } else if (self == p.getRight()) {
                        if (brother.getLeft() == null) {
                            leftRotate(brother);
                            brother = p.getLeft();
                        }
                        switchColor(brother, brother.getLeft());
                    }
                    if (self == p.getLeft()) {
                        leftRotate(p);
                    } else {
                        rightRotate(p);
                    }
                } else {  // 如果新兄弟节点无儿子
                    setColor(brother, RED);
                }
                doRemove(self);
                return;
            }

            // 情况5, 父亲是黑节点, 且兄弟是黑且无红儿子
            setColor(brother, RED);
            self = p;
        } while (self.getParent() != null);

        if (self == p.getLeft()) {
            setColor(p.getRight(), RED);
        } else {
            setColor(p.getLeft(), RED);
        }
        doRemove(_234LeaveNode);
    }

    private void switchColor(RBNode<K, V> node1, RBNode<K, V> node2) {
        if (colorOf(node1) != colorOf(node2)) {
            setColor(node2, !setColor(node1, node2.getColor()));
        }
    }

    private RBNode<K, V> getPredecessor(RBNode<K, V> node) {
        return (RBNode<K, V>) bsTree.getPredecessor(node);
    }

    private RBNode<K, V> getNode(K key) {
        return (RBNode<K, V>) bsTree.getNode(key);
    }

    /**
     * 节点左旋
     * @param p 要左旋的节点
     */
    public void leftRotate(RBNode<K, V> p) {
        if (p == null) { return; }
        RBNode<K, V> pRight = p.getRight();

        // 让右节点取代p的位置
        pRight.setParent(p.getParent());
        if (p.getParent() == null) {
            this.setRoot(pRight);
        } else if (p.getParent().getLeft() == p) {
            p.getParent().setLeft(pRight);
        } else {
            p.getParent().setRight(pRight);
        }

        // 让右节点的左节点的父节点指向p
        p.setRight(pRight.getLeft());
        if (pRight.getLeft() != null) {
            pRight.getLeft().setParent(p);
        }


        pRight.setLeft(p);
        p.setParent(pRight);
    }

    /**
     * 节点右旋
     * @param p 要右旋的节点
     */
    public void rightRotate(RBNode<K, V> p) {
        if (p == null) { return; }
        RBNode<K, V> pLeft = p.getLeft();

        // 让左节点取代p的位置
        pLeft.setParent(p.getParent());
        if (p.getParent() == null) {
            this.setRoot(pLeft);
        } else if (p.getParent().getLeft() == p) {
            p.getParent().setLeft(pLeft);
        } else {
            p.getParent().setRight(pLeft);
        }

        // 让左节点的右节点的父节点指向p
        p.setLeft(pLeft.getRight());
        if (pLeft.getRight() != null) {
            pLeft.getRight().setParent(p);
        }

        pLeft.setRight(p);
        p.setParent(pLeft);
    }

    private boolean setColor(RBNode<K, V> node, boolean color) {
        node.color = color;
        return color;
    }
    private boolean colorOf(RBNode<K, V> node) { return node == null ? BLACK : node.getColor(); }

    private void fixAfterPut(RBNode<K, V> node){
        while(node != null  && colorOf(node.getParent()) == RED) {

            RBNode<K, V> p = node.getParent();
            RBNode<K, V> pp = p.getParent();
            RBNode<K, V> uncle = pp.getLeft() == p ? pp.getRight() : pp.getLeft();

            if (colorOf(uncle) == RED) { // 情况1
                setColor(p, setColor(uncle, BLACK));
                setColor(pp, RED);
                node = node.getParent().getParent();
            } else {
                if (pp.getLeft() == p && p.getLeft() == node) {         // 左斜情况2
                    rightRotate(pp);
                    setColor(p, BLACK);
                    setColor(pp, RED);
                    node = node.getParent().getParent();
                } else if(pp.getRight() == p && p.getRight() == node) { // 右斜情况2
                    leftRotate(pp);
                    setColor(p, BLACK);
                    setColor(pp, RED);
                    node = node.getParent().getParent();
                } else if (pp.getLeft() == p && p.getRight() == node) { // 左角情况3
                    leftRotate(p);
                    rightRotate(pp);
                    setColor(node, BLACK);
                    setColor(pp, RED);
                    node = node.getParent();
                } else {                                                // 右角情况3
                    rightRotate(p);
                    leftRotate(pp);
                    setColor(node, BLACK);
                    setColor(pp, RED);
                    node = node.getParent();
                }
            }

            if (node != null && node.getParent() == null && node.getColor() == RED) { // 根节点为红-情况4
                setColor(node, BLACK);
                return;
            }

        }
    }
}
